/*-------------------------------------------------------------------------
 *
 * basebackup_zstd.c
 *	  基本备份接收器实现zstd压缩。
 *
 * Portions Copyright (c) 2010-2022, PostgreSQL Global Development Group
 *
 * IDENTIFICATION
 *	  src/backend/backup/basebackup_zstd.c
 *
 *-------------------------------------------------------------------------
 */
#include "postgres.h"

#ifdef USE_ZSTD
#include <zstd.h>
#endif

#include "backup/basebackup_sink.h"

#ifdef USE_ZSTD

typedef struct bbsink_zstd
{
	/* 所有类型的接收器的通用信息。 */
	bbsink		base;

	/* 压缩选项 */
	pg_compress_specification *compress;

	ZSTD_CCtx  *cctx;
	ZSTD_outBuffer zstd_outBuf;
} bbsink_zstd;

static void fc_bbsink_zstd_begin_backup(bbsink *fc_sink);
static void fc_bbsink_zstd_begin_archive(bbsink *fc_sink, const char *fc_archive_name);
static void fc_bbsink_zstd_archive_contents(bbsink *fc_sink, size_t fc_avail_in);
static void fc_bbsink_zstd_manifest_contents(bbsink *fc_sink, size_t fc_len);
static void fc_bbsink_zstd_end_archive(bbsink *fc_sink);
static void fc_bbsink_zstd_cleanup(bbsink *fc_sink);
static void fc_bbsink_zstd_end_backup(bbsink *fc_sink, XLogRecPtr fc_endptr,
								   TimeLineID fc_endtli);

static const bbsink_ops bbsink_zstd_ops = {
	.begin_backup = fc_bbsink_zstd_begin_backup,
	.begin_archive = fc_bbsink_zstd_begin_archive,
	.archive_contents = fc_bbsink_zstd_archive_contents,
	.end_archive = fc_bbsink_zstd_end_archive,
	.begin_manifest = bbsink_forward_begin_manifest,
	.manifest_contents = fc_bbsink_zstd_manifest_contents,
	.end_manifest = bbsink_forward_end_manifest,
	.end_backup = fc_bbsink_zstd_end_backup,
	.cleanup = fc_bbsink_zstd_cleanup
};
#endif

/*
 * 创建一个新的 basebackup sink，执行 zstd 压缩。
 */
bbsink * bbsink_zstd_new(bbsink *fc_next, pg_compress_specification *fc_compress)
{
#ifndef USE_ZSTD
	ereport(ERROR,
			(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
			 errmsg("zstd compression is not supported by this build")));
	return NULL;				/* 保持编译器安静 */
#else
	bbsink_zstd *fc_sink;

	Assert(fc_next != NULL);

	fc_sink = palloc0(sizeof(bbsink_zstd));
	*((const bbsink_ops **) &fc_sink->base.bbs_ops) = &bbsink_zstd_ops;
	fc_sink->base.bbs_next = fc_next;
	fc_sink->compress = fc_compress;

	return &fc_sink->base;
#endif
}

#ifdef USE_ZSTD

/*
 * 开始备份。
 */
static void fc_bbsink_zstd_begin_backup(bbsink *fc_sink)
{
	bbsink_zstd *fc_mysink = (bbsink_zstd *) fc_sink;
	size_t		fc_output_buffer_bound;
	size_t		fc_ret;
	pg_compress_specification *fc_compress = fc_mysink->compress;

	fc_mysink->cctx = ZSTD_createCCtx();
	if (!fc_mysink->cctx)
		elog(ERROR, "could not create zstd compression context");

	fc_ret = ZSTD_CCtx_setParameter(fc_mysink->cctx, ZSTD_c_compressionLevel,
								 fc_compress->level);
	if (ZSTD_isError(fc_ret))
		elog(ERROR, "could not set zstd compression level to %d: %s",
			 fc_compress->level, ZSTD_getErrorName(fc_ret));

	if ((fc_compress->options & PG_COMPRESSION_OPTION_WORKERS) != 0)
	{
		/*
		 * 在旧版本的 libzstd 中，此选项不存在，尝试设置它将失败。
		 * 对于较新版本，如果没有编译线程支持也是如此。
		 */
		fc_ret = ZSTD_CCtx_setParameter(fc_mysink->cctx, ZSTD_c_nbWorkers,
									 fc_compress->workers);
		if (ZSTD_isError(fc_ret))
			ereport(ERROR,
					errcode(ERRCODE_INVALID_PARAMETER_VALUE),
					errmsg("could not set compression worker count to %d: %s",
						   fc_compress->workers, ZSTD_getErrorName(fc_ret)));
	}

	/*
	 * 我们需要自己的缓冲区，因为我们将要传递给下一个接收器的数据
	 * 与传递给我们的数据不同。
	 */
	fc_mysink->base.bbs_buffer = palloc(fc_mysink->base.bbs_buffer_length);

	/*
	 * 确保下一个 sink 的 bbs_buffer 足够大，以容纳
	 * 压缩输入缓冲区。
	 */
	fc_output_buffer_bound = ZSTD_compressBound(fc_mysink->base.bbs_buffer_length);

	/*
	 * 缓冲区长度预计应该是BLCKSZ的倍数，因此向上取整。
	 */
	fc_output_buffer_bound = fc_output_buffer_bound + BLCKSZ -
		(fc_output_buffer_bound % BLCKSZ);

	bbsink_begin_backup(fc_sink->bbs_next, fc_sink->bbs_state, fc_output_buffer_bound);
}

/*
 * 准备压缩下一个归档。
 */
static void fc_bbsink_zstd_begin_archive(bbsink *fc_sink, const char *fc_archive_name)
{
	bbsink_zstd *fc_mysink = (bbsink_zstd *) fc_sink;
	char	   *fc_zstd_archive_name;

	/*
	 * 在每个档案的开始，我们重置状态以开始一个新的
	 * 压缩操作。参数是粘性的，它们会保留
	 * 在我们使用选项 ZSTD_reset_session_only 进行重置时。
	 */
	ZSTD_CCtx_reset(fc_mysink->cctx, ZSTD_reset_session_only);

	fc_mysink->zstd_outBuf.dst = fc_mysink->base.bbs_next->bbs_buffer;
	fc_mysink->zstd_outBuf.size = fc_mysink->base.bbs_next->bbs_buffer_length;
	fc_mysink->zstd_outBuf.pos = 0;

	/* 在档案名称中添加 ".zst"。 */
	fc_zstd_archive_name = psprintf("%s.zst", fc_archive_name);
	Assert(fc_sink->bbs_next != NULL);
	bbsink_begin_archive(fc_sink->bbs_next, fc_zstd_archive_name);
	pfree(fc_zstd_archive_name);
}

/*
 * 将输入数据压缩到输出缓冲区，直到我们用完输入
 * 数据。每当输出缓冲区低于输入缓冲区的压缩边界时，调用
 * 下一个 sink 的 archive_contents() 方法。
 *
 * 请注意，由于我们在压缩输入，可能很常见
 * 的情况是我们消耗所有输入数据而没有填满输出缓冲区。在
 * 这种情况下，当前输入数据的压缩表示将不会
 * 实际发送到下一个 bbsink，直到稍后调用此函数，
 * 或者甚至可能在调用 bbsink_zstd_end_archive() 之前。
 */
static void fc_bbsink_zstd_archive_contents(bbsink *fc_sink, size_t fc_len)
{
	bbsink_zstd *fc_mysink = (bbsink_zstd *) fc_sink;
	ZSTD_inBuffer fc_inBuf = {fc_mysink->base.bbs_buffer, fc_len, 0};

	while (fc_inBuf.pos < fc_inBuf.size)
	{
		size_t		fc_yet_to_flush;
		size_t		fc_max_needed = ZSTD_compressBound(fc_inBuf.size - fc_inBuf.pos);

		/*
		 * 如果输出缓冲区没有留下足够的空间，发送输出
		 * 缓冲区到下一个 sink，并重置它。
		 */
		if (fc_mysink->zstd_outBuf.size - fc_mysink->zstd_outBuf.pos < fc_max_needed)
		{
			bbsink_archive_contents(fc_mysink->base.bbs_next,
									fc_mysink->zstd_outBuf.pos);
			fc_mysink->zstd_outBuf.dst = fc_mysink->base.bbs_next->bbs_buffer;
			fc_mysink->zstd_outBuf.size =
				fc_mysink->base.bbs_next->bbs_buffer_length;
			fc_mysink->zstd_outBuf.pos = 0;
		}

		fc_yet_to_flush = ZSTD_compressStream2(fc_mysink->cctx, &fc_mysink->zstd_outBuf,
											&fc_inBuf, ZSTD_e_continue);

		if (ZSTD_isError(fc_yet_to_flush))
			elog(ERROR,
				 "could not compress data: %s",
				 ZSTD_getErrorName(fc_yet_to_flush));
	}
}

/*
 * zstd 的内部缓冲区中可能会有一些数据；我们需要将其
 * 刷新出来，同时结束 zstd 帧，然后将其转发到
 * 后续的 sink 作为档案内容。
 *
 * 然后我们可以结束此档案的处理。
 */
static void fc_bbsink_zstd_end_archive(bbsink *fc_sink)
{
	bbsink_zstd *fc_mysink = (bbsink_zstd *) fc_sink;
	size_t		fc_yet_to_flush;

	do
	{
		ZSTD_inBuffer fc_in = {NULL, 0, 0};
		size_t		fc_max_needed = ZSTD_compressBound(0);

		/*
		 * 如果输出缓冲区没有留下足够的空间，发送输出
		 * 缓冲区到下一个 sink，并重置它。
		 */
		if (fc_mysink->zstd_outBuf.size - fc_mysink->zstd_outBuf.pos < fc_max_needed)
		{
			bbsink_archive_contents(fc_mysink->base.bbs_next,
									fc_mysink->zstd_outBuf.pos);
			fc_mysink->zstd_outBuf.dst = fc_mysink->base.bbs_next->bbs_buffer;
			fc_mysink->zstd_outBuf.size =
				fc_mysink->base.bbs_next->bbs_buffer_length;
			fc_mysink->zstd_outBuf.pos = 0;
		}

		fc_yet_to_flush = ZSTD_compressStream2(fc_mysink->cctx,
											&fc_mysink->zstd_outBuf,
											&fc_in, ZSTD_e_end);

		if (ZSTD_isError(fc_yet_to_flush))
			elog(ERROR, "could not compress data: %s",
				 ZSTD_getErrorName(fc_yet_to_flush));

	} while (fc_yet_to_flush > 0);

	/* 确保将任何剩余的字节传递给下一个 sink。 */
	if (fc_mysink->zstd_outBuf.pos > 0)
		bbsink_archive_contents(fc_mysink->base.bbs_next,
								fc_mysink->zstd_outBuf.pos);

	/* 传递有关该归档已结束的信息。 */
	bbsink_forward_end_archive(fc_sink);
}

/*
 * 释放资源和上下文。
 */
static void fc_bbsink_zstd_end_backup(bbsink *fc_sink, XLogRecPtr fc_endptr,
					   TimeLineID fc_endtli)
{
	bbsink_zstd *fc_mysink = (bbsink_zstd *) fc_sink;

	/* 释放上下文。 */
	if (fc_mysink->cctx)
	{
		ZSTD_freeCCtx(fc_mysink->cctx);
		fc_mysink->cctx = NULL;
	}

	bbsink_forward_end_backup(fc_sink, fc_endptr, fc_endtli);
}

/*
 * 清单内容未压缩，但我们确实需要将其复制到
 * 后续接收器的缓冲区中，因为我们有自己的缓冲区。
 */
static void fc_bbsink_zstd_manifest_contents(bbsink *fc_sink, size_t fc_len)
{
	memcpy(fc_sink->bbs_next->bbs_buffer, fc_sink->bbs_buffer, fc_len);
	bbsink_manifest_contents(fc_sink->bbs_next, fc_len);
}

/*
 * 如果备份失败，确保释放任何分配的压缩上下文，
 * 以免导致内存泄漏。
 */
static void fc_bbsink_zstd_cleanup(bbsink *fc_sink)
{
	bbsink_zstd *fc_mysink = (bbsink_zstd *) fc_sink;

	/* 如果尚未释放，则释放上下文。 */
	if (fc_mysink->cctx)
	{
		ZSTD_freeCCtx(fc_mysink->cctx);
		fc_mysink->cctx = NULL;
	}
}

#endif
